import graphviz

# 初始化有向图
dot = graphviz.Digraph(comment='Object Detection Timeline',
                      format='svg',
                      engine='dot',
                      graph_attr={'bgcolor': 'transparent'})

# 全局样式设置
dot.attr('node', 
        shape='rect',
        style='filled',
        fontname='Arial',
        fontsize='10',
        fillcolor='#F0F0F0',
        color='#333333')
dot.attr('edge', 
        color='#666666',
        arrowsize='0.7')

# ================= 时间轴主线 =================
with dot.subgraph(name='timeline') as tl:
    tl.attr(rank='same')
    years = [2001, 2006, 2008, 2013, 2015, 2016, 2017, 
            2018, 2019, 2020, 2022, 2023, 2025]
    
    # 创建透明节点形成时间轴
    for year in years:
        tl.node(f'y{year}', 
               label=f'|{year}|',
               shape='plaintext',
               width='0.2')

# ============ 传统方法分支 ============
with dot.subgraph(name='traditional') as trad:
    trad.attr(color='#4B8BBE')
    models = [
        ('vj', 'Viola-Jones\n(2001)', 2001),
        ('hog', 'HOG\n(2006)', 2006),
        ('dpm', 'DPM\n(2008)', 2008)
    ]
    for i, (id, label, year) in enumerate(models):
        trad.node(id, label, fillcolor='#8FAADC')
        if i > 0:
            trad.edge(models[i-1][0], id)

# ============ 深度学习分叉点 ============
dot.node('dl_start', 'Deep Learning\nEra',
        shape='diamond',
        fillcolor='#FFD43B',
        fontsize='12')

# 连接传统到深度学习
dot.edge('dpm', 'dl_start', 
        lhead='cluster_dl',
        constraint='false')

# ============ 单阶段检测器分支 ============
with dot.subgraph(name='cluster_onestage') as one_stage:
    one_stage.attr(label='Single-Stage Detectors',
                  color='#72B352',
                  labelloc='t')
    models = [
        ('overfeat', 'OverFeat\n(2013)', 2013),
        ('yolo', 'YOLO\n(2015)', 2015),
        ('ssd', 'SSD\n(2016)', 2016),
        ('yolov2', 'YOLOv2\n(2017)', 2017),
        ('retinanet', 'RetinaNet\n(2018)', 2018),
        ('yolov3', 'YOLOv3\n(2018)', 2018),
        ('yolov4', 'YOLOv4\n(2020)', 2020),
        ('yolov5', 'YOLOv5\n(2020)', 2020),
        ('yolov6', 'YOLOv6\n(2022)', 2022),
        ('yolov7', 'YOLOv7\n(2023)', 2023),
        ('yolo2025', 'YOLO-X\n(2025?)', 2025)
    ]
    for i, (id, label, year) in enumerate(models):
        one_stage.node(id, label, fillcolor='#A9D08E')
        if i > 0:
            one_stage.edge(models[i-1][0], id)

# ============ 二阶段检测器分支 ============
with dot.subgraph(name='cluster_twostage') as two_stage:
    two_stage.attr(label='Two-Stage Detectors',
                  color='#C54B6C',
                  labelloc='t')
    models = [
        ('rcnn', 'R-CNN\n(2013)', 2013),
        ('fast', 'Fast R-CNN\n(2015)', 2015),
        ('faster', 'Faster R-CNN\n(2016)', 2016),
        ('mask', 'Mask R-CNN\n(2017)', 2017),
        ('cascade', 'Cascade R-CNN\n(2018)', 2018),
        ('dynamic', 'Dynamic R-CNN\n(2020)', 2020),
        ('smarter', 'Smarter R-CNN\n(2025?)', 2025)
    ]
    for i, (id, label, year) in enumerate(models):
        two_stage.node(id, label, fillcolor='#E6B8B7')
        if i > 0:
            two_stage.edge(models[i-1][0], id)

# ============ 连接分叉点 ============
dot.edge('dl_start', 'overfeat', lhead='cluster_onestage')
dot.edge('dl_start', 'rcnn', lhead='cluster_twostage')

# ============ 未来技术预测 ============
dot.node('transformer', 'Transformer-Based\nDetectors (2022+)',
        shape='ellipse',
        fillcolor='#9467BD')
dot.edge('yolov7', 'transformer')
dot.edge('smarter', 'transformer')

# ============ 生成输出 ============
dot.attr(label='Object Detection Development Timeline\n(2001-2025 Projection)',
        labelloc='t',
        fontsize='16',
        fontname='Arial Bold')
dot.render('detection_timeline', cleanup=True, view=True)